import matplotlib
import matplotlib.pyplot as plt
import numpy as np

from tensorflow.keras.callbacks import History


def pic_it(*, epochs: int, record: History,
           metric: str = "acc", other: bool = False,
           file_path: str = None):
    """
    画图函数

    :param epochs: 训练次数
    :param record: 训练后的记录数据
    :param metric:
    :param other: 是否启用
    :param file_path: 与other一起使用，文件保存位置
    :return:
    """
    # 画图
    if other:
        matplotlib.use("Agg")
    plt.style.use("ggplot")
    plt.figure()
    plt.plot(np.arange(0, epochs), record.history["loss"], label="train_loss")
    plt.plot(np.arange(0, epochs), record.history["val_loss"], label="val_loss")
    plt.plot(np.arange(0, epochs), record.history[f"{metric}"], label=f"train_{metric}")
    plt.plot(np.arange(0, epochs), record.history[f"val_{metric}"], label=f"val_{metric}")
    plt.title("Training Loss and Accuracy")
    plt.xlabel("Epoch #")
    plt.ylabel("Loss/Accuracy")
    plt.legend()
    if file_path:
        plt.savefig(file_path)
    else:
        plt.show()
